package com.xiashitech.agent.instrumentation.base.context;

import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.baggage.BaggageEntry;

import javax.annotation.Nullable;
import javax.annotation.concurrent.Immutable;
import java.util.Map;
import java.util.function.BiConsumer;

@Immutable
public interface XSBaggage extends XSImplicitContextKeyed {
    public static final ThreadLocal<String> userIdThreadLocal = new ThreadLocal();

    static XSBaggage empty() {
        return XSImmutableBaggage.empty();
    }

    static XSBaggageBuilder builder() {
        return XSImmutableBaggage.builder();
    }

    static XSBaggage[] currentWithConfluence() {
        XSBaggage currentXSBaggage = current();
        XSBaggage[] confluenceXSBaggages = confluence();
        XSBaggage[] xsBaggages = new XSBaggage[confluenceXSBaggages.length + 1];
        if(confluenceXSBaggages.length == 0) {
            xsBaggages[0] = currentXSBaggage;
        } else {
            xsBaggages[0] = currentXSBaggage;
            System.arraycopy(confluenceXSBaggages, 0, xsBaggages, 1, confluenceXSBaggages.length);
        }
        return xsBaggages;
    }

    static XSBaggage current() {
        return fromContext(XSContext.current());
    }

    /**
     * 汇合各个下游分支的baggage
     * @return
     */
    static XSBaggage[] confluence() {
        return fromContextForConfluence(XSContext.current());
    }

    static XSBaggage[] fromContextForConfluence(XSContext context) {
        XSBaggage[] xsBaggages = context.get(XSBaggageContextKey.ARRAY_KEY);
        return xsBaggages != null ? xsBaggages : new XSBaggage[]{};
    }

    static XSBaggage fromContext(XSContext context) {
        XSBaggage baggage = context.get(XSBaggageContextKey.KEY);
        return baggage != null ? baggage : empty();
    }

    @Nullable
    static XSBaggage fromContextOrNull(XSContext context) {
        return context.get(XSBaggageContextKey.KEY);
    }

    default XSContext storeInContext(XSContext context) {
        return context.with(XSBaggageContextKey.KEY, this);
    }

    /**
     * 将当天baggage存储在上下中汇合后的baggage数组中
     * @param context
     * @return
     */
    default XSContext storeInContextForConfluence(XSContext context) {
        XSBaggage[] xsBaggages = context.get(XSBaggageContextKey.ARRAY_KEY);
        if(xsBaggages != null) {
            if(xsBaggages.length == 0) {
                return context.with(XSBaggageContextKey.ARRAY_KEY, new XSBaggage[]{this});
            }
            boolean isExist = false;
            for (int i = 0; i < xsBaggages.length; i++) {
                if(xsBaggages[i] == this || xsBaggages[i].equals(this)) {
                    isExist = true;
                    break;
                }
            }
            if(!isExist) {
                XSBaggage[] newXSBaggages = new XSBaggage[xsBaggages.length + 1];
                System.arraycopy(xsBaggages, 0, newXSBaggages, 0, xsBaggages.length);
                newXSBaggages[xsBaggages.length - 1] = this;
                return context.with(XSBaggageContextKey.ARRAY_KEY, newXSBaggages);
            }
            return context;
        } else {
            return context.with(XSBaggageContextKey.ARRAY_KEY, new XSBaggage[]{this});
        }
    }

    int size();

    default boolean isEmpty() {
        return this.size() == 0;
    }

    void forEach(BiConsumer<? super String, ? super BaggageEntry> consumer);

    Map<String, BaggageEntry> asMap();

    @Nullable
    String getEntryValue(String key);

    XSBaggageBuilder toBuilder();
}
